﻿#pragma once

#include <iostream>
#include <vector>
#include <iomanip>
#include <unordered_map>
#include <optional>

using namespace std;

template <class V, class W, bool IsDirected = false>
class Graph {
public:
    Graph(const vector<V> &vertices) : _vertices(vertices) {
        int n = _vertices.size();
        _edges.resize(n, vector<optional<W>>(n, nullopt));

        for (int i = 0; i < n; ++i) {
            _vertexToIndex.insert({ _vertices[i], i });
            _edges[i][i] = 0; // 顶点到自身的距离为0
        }
    }

    bool addEdge(const V& src, const V& dst, const W& weight) {
        int srcIdx = getVertexIndex(src);
        int dstIdx = getVertexIndex(dst);

        if (srcIdx == -1 || dstIdx == -1) {
            return false;
        }

        _edges[srcIdx][dstIdx] = weight;
        if (!IsDirected) {
            _edges[dstIdx][srcIdx] = weight;
        }

        return true;
    }

    int getVertexIndex(const V& v) const {
        auto it = _vertexToIndex.find(v);
        return it != _vertexToIndex.end() ? it->second : -1;
    }

    void print() const {
        // 打印列标题
        cout << "     ";
        for (const auto& v : _vertices) {
            cout << setw(5) << v;
        }
        cout << endl;

        // 打印矩阵内容
        for (size_t i = 0; i < _edges.size(); ++i) {
            cout << setw(5) << _vertices[i];
            for (const auto &val : _edges[i]) {
                if (val) {
                    cout << setw(5) << *val;
                }
                else {
                    cout << setw(5) << "INF";
                }
            }
            cout << endl;
        }
    }

private:
    vector<V> _vertices;
    unordered_map<V, int> _vertexToIndex;
    vector<vector<optional<W>>> _edges;
};